import subprocess
import psutil
import socket
import re
import platform

##############
#author:张子路#
##############

class LinuxCheck():
    def __init__(self):
        pass

    def Ckeck(self,CheckList):
        try:        #条目核查
            res = []
            for iteam in CheckList:
                val = self.run_code(iteam[1]).strip()
                if val == '' or val != iteam[2]:
                    res.append([iteam[0],'设置不合格',iteam[3]])
                else:
                    res.append([iteam[0],'合格','保持当前设置。'])
            return res
        except Exception:
            print('条目核查失败。')


    def item_check(self):
        Check_res = []
        CheckList=[['DHCP Sever','systemctl is-enabled isc-dhcp-server','disabled','建议将DHCP Sever服务设置为disabled'],
                   ['accept_source_route',"sysctl net.ipv4.conf.all.accept_source_route | awk '{print $3}'",'0','建议将net.ipv4.conf.all.accept_source_route设置为0'],
                   ['SSH-X11Forwarding',r"grep '^X11Forwarding' /etc/ssh/sshd_config | awk '{print $2}'",'no','建议将X11Forwarding项的值置为no']]
        Check_res=self.Ckeck(CheckList)
        return Check_res

    def get_IP(self):
        try:        #获取IP
            net_info = self.run_code('ifconfig')
            reg = re.compile(r'inet ((?:\d{1,3}\.){3}\d{1,3})')
            ip_list = []
            for m in reg.finditer(net_info):
                ip_list.append(m.group(1))
            # print('ip list:', ip_list)
            return ip_list
        except Exception:
            print('获取IP失败。')


    def get_ports2(self):
        try:    #端口信息查询
            info = self.run_code('netstat -ano | grep LISTEN')
            reg = re.compile(r'tcp\s+\d+\s+\d+\s(?:\d{1,3}\.){3}\d{1,3}:(\d+)')
            ports = []
            for m in reg.finditer(info):
                if m.group(1) == '0':
                    continue
                ports.append(m.group(1))
            return ports
        except Exception:
            print('端口信息查询失败。')


    def get_mem_info(self):
        try:       #内存核查失败
            pc_mem = psutil.virtual_memory()
            G = 1024**3
            return {
                'total_GB': pc_mem.total/G,
                'used_GB': pc_mem.used/G,
                'percent': pc_mem.used/pc_mem.total * 100,
            }
        except Exception:
            print('内存核查失败。')


    def get_host_info(self) -> dict:
        try:         #主机信息核查
            info_res = {}
            hostname = socket.gethostname()
            info_res['hostname'] = hostname
            info_res['ip'] = self.get_IP()
            info_res['OS'] = platform.platform()
            info_res['CPU'] = {
                'model': platform.uname().processor,
                'percent': psutil.cpu_percent()
            }
            info_res['opening_port'] = self.get_ports2()
            info_res['mem'] = self.get_mem_info()
            return info_res
        except Exception:
            print('主机信息核查出错。')

    def output_dict(self,d: dict):
        try:
            for key in d.keys():         #输出结果
                print(key, ':', d[key])
        except Exception:
            print("输出出错。")

    def output_list(self,CheckList):
        try:
            for iteam in CheckList:     #输出结果
                print("核查条目:"+iteam[0]+'\n核查结果:'+iteam[1]+'\n加固建议:'+iteam[2]+'\n\n')
        except Exception:
            print("输出出错。")


    def Start(self):
        try:        #开始检查
            print('主机信息:')
            self.output_dict(self.get_host_info())
            print("条目核查结果:")
            self.output_list(self.item_check())
            str='{ if($3 == 0 && $1 != "root") print $1 }'
            re=self.run_code("awk -F: '{0}' /etc/passwd ".format(str)).strip()
            if re==None or re=='':
                print('\n后门核查：\n存在的UID为0，不为root的用户:\n'+'无')
            else:
                print('\n后门核查：\n存在的UID为0，不为root的用户:'+self.run_code("awk -F: '{0}' /etc/passwd ".format(str)))
        except Exception:
            print("核查出错。")

    def run_code(self,cmd):
        try:         #运行命令
            res = subprocess.Popen(
                cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            for line in res.communicate():
                    return line.decode('utf-8')
        except Exception:
            print("运行命令出错。")

if __name__ == '__main__':
    LinuxCheckClass=LinuxCheck()
    LinuxCheckClass.Start()


